package com.sql.parser.lineage.doris;


import doris.DorisParser;
import doris.DorisParserBaseListener;
import org.antlr.v4.runtime.Token;
import org.antlr.v4.runtime.TokenStream;
import org.antlr.v4.runtime.TokenStreamRewriter;

import java.util.HashMap;


public class SqlListenerRewrite extends DorisParserBaseListener {
    String dbName;
    HashMap<String, String> sourceTableMapWith = new HashMap<>();
    HashMap<Token, String> tokenHash = new HashMap<>();
    String sqlStr;
    public TokenStreamRewriter rewriter;


    public SqlListenerRewrite(TokenStream tokens) {
        rewriter = new TokenStreamRewriter(tokens);
    }

    public SqlListenerRewrite(TokenStream tokens, String dbName) {
        rewriter = new TokenStreamRewriter(tokens);
        this.dbName = dbName;
    }

    @Override
    public void enterSingleStatement(DorisParser.SingleStatementContext ctx) {
        sqlStr = ctx.getText();
    }


    @Override
    public void enterTableName(DorisParser.TableNameContext ctx) {
        String curTableName = ctx.multipartIdentifier().getText();
        if (!curTableName.contains(".")) {
            String replaceDbName;
            if (curTableName.isEmpty()) {
                replaceDbName = "";
            } else {
                replaceDbName = dbName + "." + curTableName;
            }
            tokenHash.put(ctx.getStart(), replaceDbName);
        }
    }

    @Override
    public void exitSelectClause(DorisParser.SelectClauseContext ctx) {
        rw();
    }

    @Override
    public void exitSingleStatement(DorisParser.SingleStatementContext ctx) {
        rw();
        super.exitSingleStatement(ctx);
    }

    public void rw() {
        for (Token key : tokenHash.keySet()) {
            rewriter.replace(key, tokenHash.get(key));
        }
    }
}
